import torch
import torch.nn as nn


class FocalLoss(nn.Module):
    def __init__(self, weight=None, reduction='mean', gamma=0, eps=1e-7):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.eps = eps
        self.ce = torch.nn.CrossEntropyLoss(weight=weight, reduction=reduction)

    def forward(self, inputs, targets):
        # input shape: (batch_size, num_logits)
        # target shape: (batch, )
        logp = self.ce(inputs, targets)
        p = torch.exp(-logp)
        loss = -(1 - p) ** self.gamma * logp
        return torch.mean(loss)




